/**
 * Copyright (C) 2010 - 2013 Harry Glasgow
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.googlecode.jaden.engine;

import com.googlecode.jaden.common.config.ProjectConfiguration;
import com.googlecode.jaden.common.other.JadenException;
import com.googlecode.jaden.common.other.SmartTokenizer;
import com.googlecode.jaden.common.other.Utils;
import com.googlecode.jaden.ui.Constants;
import org.apache.poi.hssf.usermodel.HSSFCell;
import org.apache.poi.hssf.usermodel.HSSFRow;
import org.apache.poi.hssf.usermodel.HSSFSheet;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;

import java.io.*;
import java.lang.reflect.Array;
import java.util.*;

public class DataSet {

    /**
     * Seed for repeatable randomness.
     */
    public static final int SEED = 42;

    private String[][] trainingInputValues; // [row][column]
    private String[] trainingOutputValues; // [row]
    private String[][] validationInputValues; // [row][column]
    private String[] validationOutputValues; // [row]
    private String[][] testingInputValues; // [row][column]
    private String[] testingOutputValues; // [row]

    public DataSet(ProjectConfiguration configuration) throws JadenException {

        Map<Integer, String> trainingFileLines = new TreeMap<Integer, String>();
        Map<Integer, String> validationFileLines = new TreeMap<Integer, String>();
        Map<Integer, String> testingFileLines = new TreeMap<Integer, String>();

        // Get training file.
        String columnSeparator = configuration.getColumnSeparator();
        loadFile(configuration.getTrainingFilePath(), trainingFileLines, columnSeparator);

        // Get validation file.
        boolean validationFromTrainingFile = configuration.isValidationFromTrainingFile();
        if (!validationFromTrainingFile) {
            loadFile(configuration.getValidationFilePath(), validationFileLines, columnSeparator);
        }

        // Get testing file.
        boolean testingFromTrainingFile = configuration.isTestingFromTrainingFile();
        if (!testingFromTrainingFile) {
            loadFile(configuration.getTestingFilePath(), testingFileLines, columnSeparator);
        }

        Map<Integer, String[]> parsedTraining = new TreeMap<Integer, String[]>();
        int dataLength = 0;
        boolean gotFirstLine = false;
        String stringLimiter = configuration.isUseStringLimiter() ? configuration.getStringLimiter() : "";
        boolean scanToEnd = configuration.isScanToEnd();
        int firstLine = configuration.getFirstLine();
        int lastLine = configuration.getLastLine();
        int[] inputColumnMappings = configuration.getInputColumnMappings();
        int outputColumnMapping = configuration.getOutputColumnMapping();
        boolean coalesceDuplicateTokens = configuration.isCoalesceDuplicateColumnSeparators();
        for (Map.Entry<Integer, String> entry : trainingFileLines.entrySet()) {
            int lineNumber = entry.getKey();
            String line = entry.getValue();
            if (lineNumber >= firstLine && (scanToEnd || lineNumber <= lastLine)) {
                String[] strings = parse(line, columnSeparator, stringLimiter, coalesceDuplicateTokens);
                if (!gotFirstLine) {
                    gotFirstLine = true;
                    dataLength = strings.length;
                    if (!validateDataLength(dataLength, inputColumnMappings, outputColumnMapping)) {
                        throw new JadenException(Utils.translate(Constants.ERROR_INCONSISTENT_TRAINING_FILE_COLUMNS, String.valueOf(dataLength)));
                    }
                } else if (dataLength != strings.length) {
                    throw new JadenException(Utils.translate(Constants.ERROR_TRAINING_FILE_DIFFERENT_LENGTHS, String.valueOf(lineNumber)));
                }
                parsedTraining.put(lineNumber, strings);
            }
        }

        Map<Integer, String[]> parsedValidation = new TreeMap<Integer, String[]>();
        if (!validationFromTrainingFile) {
            gotFirstLine = false;
            for (Map.Entry<Integer, String> entry : validationFileLines.entrySet()) {
                int lineNumber = entry.getKey();
                String line = entry.getValue();
                if (lineNumber >= firstLine && (scanToEnd || lineNumber <= lastLine)) {
                    String[] strings = parse(line, columnSeparator, stringLimiter, coalesceDuplicateTokens);
                    if (!gotFirstLine) {
                        gotFirstLine = true;
                        dataLength = strings.length;
                        if (!validateDataLength(dataLength, inputColumnMappings, outputColumnMapping)) {
                            throw new JadenException(Utils.translate(Constants.ERROR_INCONSISTENT_VALIDATION_FILE_COLUMNS, String.valueOf(dataLength)));
                        }
                    } else if (dataLength != strings.length) {
                        throw new JadenException(Utils.translate(Constants.ERROR_VALIDATION_FILE_DIFFERENT_LENGTHS, String.valueOf(lineNumber)));
                    }
                    parsedValidation.put(lineNumber, strings);
                }
            }
        }

        Map<Integer, String[]> parsedTesting = new TreeMap<Integer, String[]>();
        if (!testingFromTrainingFile) {
            gotFirstLine = false;
            for (Map.Entry<Integer, String> entry : testingFileLines.entrySet()) {
                int lineNumber = entry.getKey();
                String line = entry.getValue();
                if (lineNumber >= firstLine && (scanToEnd || lineNumber <= lastLine)) {
                    String[] strings = parse(line, columnSeparator, stringLimiter, coalesceDuplicateTokens);
                    if (!gotFirstLine) {
                        gotFirstLine = true;
                        dataLength = strings.length;
                        if (!validateDataLength(dataLength, inputColumnMappings, outputColumnMapping)) {
                            throw new JadenException(Utils.translate(Constants.ERROR_INCONSISTENT_TESTING_FILE_COLUMNS, String.valueOf(dataLength)));
                        }
                    } else if (dataLength != strings.length) {
                        throw new JadenException(Utils.translate(Constants.ERROR_TESTING_FILE_DIFFERENT_LENGTHS, String.valueOf(lineNumber)));
                    }
                    parsedTesting.put(lineNumber, strings);
                }
            }
        }

        int validationDataPercentage = configuration.getValidationDataPercentage();
        int testingDataPercentage = configuration.getTestingDataPercentage();
        if (validationFromTrainingFile && testingFromTrainingFile) {
            validationAndTestingFromFile(parsedTraining, validationDataPercentage, testingDataPercentage, inputColumnMappings,
                    outputColumnMapping);
        } else if (validationFromTrainingFile && !testingFromTrainingFile) {
            validationOnlyFromFile(parsedTraining, parsedTesting, validationDataPercentage, inputColumnMappings, outputColumnMapping);
        } else if (testingFromTrainingFile) {
            testingOnlyFromFile(parsedTraining, parsedValidation, testingDataPercentage, inputColumnMappings, outputColumnMapping);
        } else {
            individualFiles(parsedTraining, parsedValidation, parsedTesting, inputColumnMappings, outputColumnMapping);
        }

        boolean randomizeData = configuration.isRandomizeData();
        if (randomizeData) {
            randomize(trainingInputValues, trainingOutputValues);
            randomize(validationInputValues, validationOutputValues);
            randomize(testingInputValues, testingOutputValues);
        }
    }

    /**
     * Randomize the data in the training, validation and testing arrays, keeping them in sync with each other.
     *
     * @param inputValues
     * @param outputValues
     */
    private void randomize(String[][] inputValues, String[] outputValues) {
        Random random = new Random(SEED);
        int size = outputValues.length;
        String[][] tempInputValues = new String[inputValues.length][inputValues[0].length];
        System.arraycopy(inputValues, 0, tempInputValues, 0, size);
        String[] tempOutputValues = new String[outputValues.length];
        System.arraycopy(outputValues, 0, tempOutputValues, 0, size);
        boolean[] done = new boolean[size];

        for (int i = 0; i < size; i++) {
            int randomPosition;
            do {
                randomPosition = random.nextInt(size);
            } while (done[randomPosition]);
            done[randomPosition] = true;
            inputValues[i] = tempInputValues[randomPosition];
            outputValues[i] = tempOutputValues[randomPosition];
        }

    }

    /**
     * Make sure that the input and output column mappings will work with this file data line length.
     *
     * @param dataLength
     * @param inputColumnMappings
     * @param outputColumnMapping
     * @return
     */
    private boolean validateDataLength(int dataLength, int[] inputColumnMappings, int outputColumnMapping) {
        for (int inputColumnMapping : inputColumnMappings) {
            if (inputColumnMapping >= dataLength) {
                return false;
            }
        }
        return !(outputColumnMapping >= dataLength);
    }

    private static void loadFile(String filePath, Map<Integer, String> fileLines,
                                 String columnSeparator) throws JadenException {
        try {
            if (filePath.endsWith(".xls")) {
                FileInputStream fileInputStream = new FileInputStream(new File(filePath));
                HSSFWorkbook workbook = new HSSFWorkbook(fileInputStream);
                HSSFSheet worksheet = workbook.getSheetAt(0);
                int firstRowNum = worksheet.getFirstRowNum();
                int lastRowNum = worksheet.getLastRowNum();
                int lineNumber = 1;
                for (int rowNum = firstRowNum; rowNum <= lastRowNum; rowNum++) {
                    StringBuilder sb = new StringBuilder();
                    HSSFRow row = worksheet.getRow(rowNum);
                    int firstCellNum = row.getFirstCellNum();
                    int lastCellNum = row.getLastCellNum();
                    for (int cellNum = firstCellNum; cellNum <= lastCellNum; cellNum++) {
                        HSSFCell cell = row.getCell(cellNum);
                        if (cell != null) {
                            if (cell.getCellType() == HSSFCell.CELL_TYPE_STRING) {
                                String value = cell.getStringCellValue();
                                sb.append(value);
                                sb.append(columnSeparator);
                            } else if (cell.getCellType() == HSSFCell.CELL_TYPE_NUMERIC) {
                                double value = cell.getNumericCellValue();
                                sb.append(value);
                                sb.append(columnSeparator);
                            }
                        }
                    }
                    String line = sb.toString();
                    if (line.endsWith(",")) {
                        line = line.substring(0, line.length() - 1);
                    }
                    fileLines.put(lineNumber++, line);
                }
            } else {
                BufferedReader br = null;
                try {
                    br = new BufferedReader(new FileReader(new File(filePath)));
                    String line;
                    int lineNumber = 1;
                    while ((line = br.readLine()) != null) {
                        fileLines.put(lineNumber++, line);
                    }
                } finally {
                    if (br != null) {
                        try {
                            br.close();
                        } catch (IOException e) {
                            // Ignore
                        }
                    }
                }
            }
        } catch (IOException e) {
            Utils.handleException(e);
            throw new JadenException("There was a problem reading file " + filePath + ".");
        }

        // Okay, do we have one or more blanks as the last lines?
        // If so, remove them.
        int key;
        while ((key = findLastBlankKey(fileLines)) != 0) {
            fileLines.remove(key);
        }
    }

    /**
     * Find the last line # if blank, or zero.
     * @param fileLines
     * @return
     */
    private static int findLastBlankKey(Map<Integer, String> fileLines) {
        int candidate = 0;
        for (Integer key : fileLines.keySet()) {
            if (fileLines.get(key).trim().length() > 0) {
                candidate = 0;
            } else if (key > candidate) {
                candidate = key;
            }
        }
        return candidate;
    }

    private void individualFiles(Map<Integer, String[]> parsedTraining, Map<Integer, String[]> parsedValidation,
                                 Map<Integer, String[]> parsedTesting, int[] inputColumnMappings,
                                 int outputColumnMapping) {
        trainingInputValues = new String[parsedTraining.size()][inputColumnMappings.length];
        trainingOutputValues = new String[parsedTraining.size()];
        int i = 0;
        for (String[] strings : parsedTraining.values()) {
            mapInputs(strings, trainingInputValues[i], inputColumnMappings);
            trainingOutputValues[i] = mapOutputs(strings, outputColumnMapping);
            i++;
        }

        validationInputValues = new String[parsedValidation.size()][inputColumnMappings.length];
        validationOutputValues = new String[parsedValidation.size()];
        i = 0;
        for (String[] strings : parsedValidation.values()) {
            mapInputs(strings, validationInputValues[i], inputColumnMappings);
            validationOutputValues[i] = mapOutputs(strings, outputColumnMapping);
            i++;
        }

        testingInputValues = new String[parsedTesting.size()][inputColumnMappings.length];
        testingOutputValues = new String[parsedTesting.size()];
        i = 0;
        for (String[] strings : parsedTesting.values()) {
            mapInputs(strings, testingInputValues[i], inputColumnMappings);
            testingOutputValues[i] = mapOutputs(strings, outputColumnMapping);
            i++;
        }
    }

    private void testingOnlyFromFile(Map<Integer, String[]> parsedTraining, Map<Integer, String[]> parsedValidation,
                                     int testingDataPercentage, int[] inputColumnMappings, int outputColumnMapping)
            throws JadenException {
        int totalCount = parsedTraining.size();
        int testingCount = (int) (totalCount * testingDataPercentage / 100.0);
        if (testingCount == 0) {
            testingCount = 1;
        }
        int trainingCount = totalCount - testingCount;
        if (trainingCount <= 0) {
            throw new IllegalStateException("Invalid training count " + trainingCount);
        }

        trainingInputValues = new String[trainingCount][inputColumnMappings.length];
        trainingOutputValues = new String[trainingCount];
        testingInputValues = new String[testingCount][inputColumnMappings.length];
        testingOutputValues = new String[testingCount];

        int i = 0;
        int trainingPointer = 0;
        int testingPointer = 0;

        for (String[] strings : parsedTraining.values()) {
            if (i < trainingCount) {
                mapInputs(strings, trainingInputValues[trainingPointer], inputColumnMappings);
                trainingOutputValues[trainingPointer++] = mapOutputs(strings, outputColumnMapping);
            } else if (i < trainingCount + testingCount) {
                mapInputs(strings, testingInputValues[testingPointer], inputColumnMappings);
                validationOutputValues[testingPointer++] = mapOutputs(strings, outputColumnMapping);
            }
            i++;
        }

        validationInputValues = new String[parsedValidation.size()][inputColumnMappings.length];
        validationOutputValues = new String[parsedValidation.size()];

        i = 0;
        for (String[] strings : parsedValidation.values()) {
            mapInputs(strings, validationInputValues[i], inputColumnMappings);
            validationOutputValues[i] = mapOutputs(strings, outputColumnMapping);
            i++;
        }
    }

    private void validationOnlyFromFile(Map<Integer, String[]> parsedTraining, Map<Integer, String[]> parsedTesting,
                                        int validationDataPercentage, int[] inputColumnMappings, int outputColumnMapping)
            throws JadenException {
        int totalCount = parsedTraining.size();
        int validationCount = (int) (totalCount * validationDataPercentage / 100.0);
        if (validationCount == 0) {
            validationCount = 1;
        }
        int trainingCount = totalCount - validationCount;
        if (trainingCount <= 0) {
            throw new IllegalStateException("Invalid training count " + trainingCount);
        }

        trainingInputValues = new String[trainingCount][inputColumnMappings.length];
        trainingOutputValues = new String[trainingCount];
        validationInputValues = new String[validationCount][inputColumnMappings.length];
        validationOutputValues = new String[validationCount];

        int i = 0;
        int trainingPointer = 0;
        int validationPointer = 0;

        for (String[] strings : parsedTraining.values()) {
            if (i < trainingCount) {
                mapInputs(strings, trainingInputValues[trainingPointer], inputColumnMappings);
                trainingOutputValues[trainingPointer++] = mapOutputs(strings, outputColumnMapping);
            } else if (i < trainingCount + validationCount) {
                mapInputs(strings, validationInputValues[validationPointer], inputColumnMappings);
                validationOutputValues[validationPointer++] = mapOutputs(strings, outputColumnMapping);
            }
            i++;
        }

        testingInputValues = new String[parsedTesting.size()][inputColumnMappings.length];
        testingOutputValues = new String[parsedTesting.size()];

        i = 0;
        for (String[] strings : parsedTesting.values()) {
            mapInputs(strings, testingInputValues[i], inputColumnMappings);
            testingOutputValues[i] = mapOutputs(strings, outputColumnMapping);
            i++;
        }
    }

    private void validationAndTestingFromFile(Map<Integer, String[]> parsedTraining, int validationDataPercentage,
                                              int testingDataPercentage, int[] inputColumnMappings,
                                              int outputColumnMapping)
            throws JadenException {
        int dataCount = parsedTraining.size();
        int validationCount = (int) (dataCount * validationDataPercentage / 100.0);
        if (validationCount == 0) {
            validationCount = 1;
        }
        int testingCount = (int) (dataCount * testingDataPercentage / 100.0);
        if (testingCount == 0) {
            testingCount = 1;
        }
        int trainingCount = dataCount - validationCount - testingCount;
        if (trainingCount <= 0) {
            throw new IllegalStateException("Invalid training count " + trainingCount);
        }

        trainingInputValues = new String[trainingCount][inputColumnMappings.length];
        trainingOutputValues = new String[trainingCount];
        validationInputValues = new String[validationCount][inputColumnMappings.length];
        validationOutputValues = new String[validationCount];
        testingInputValues = new String[testingCount][inputColumnMappings.length];
        testingOutputValues = new String[testingCount];

        int i = 0;
        int trainingPointer = 0;
        int validationPointer = 0;
        int testingPointer = 0;

        for (String[] strings : parsedTraining.values()) {
            if (i < trainingCount) {
                mapInputs(strings, trainingInputValues[trainingPointer],
                        inputColumnMappings);
                trainingOutputValues[trainingPointer++] = mapOutputs(strings,
                        outputColumnMapping);
            } else if (i < trainingCount + validationCount) {
                mapInputs(strings, validationInputValues[validationPointer],
                        inputColumnMappings);
                validationOutputValues[validationPointer++] = mapOutputs(strings,
                        outputColumnMapping);
            } else {
                mapInputs(strings, testingInputValues[testingPointer],
                        inputColumnMappings);
                testingOutputValues[testingPointer++] = mapOutputs(strings,
                        outputColumnMapping);
            }
            i++;
        }
    }

    private static void mapInputs(String[] strings, String[] inputValue, int[] inputMappings) {
        int i = 0;
        for (int inputMapping : inputMappings) {
            inputValue[i++] = strings[inputMapping];
        }
    }

    private static String mapOutputs(String[] strings, int outputMapping) {
        return strings[outputMapping];
    }

    private static String[] parse(String line, String columnSeparator, String stringLimiter,
                                  boolean coalesceDuplicateTokens) {
        if (line.contains(columnSeparator)) {
            SmartTokenizer st = new SmartTokenizer(line, columnSeparator, stringLimiter, coalesceDuplicateTokens);
            int count = st.getParts().length;
            String[] results = new String[count];
            for (int i = 0; i < count; i++) {
                results[i] = st.getParts()[i];
            }
            return results;
        } else {
            return new String[]{line};
        }
    }

    public String[][] getTrainingInputValues() {
        return trainingInputValues;
    }

    public String[] getTrainingOutputValues() {
        return trainingOutputValues;
    }

    public String[][] getValidationInputValues() {
        return validationInputValues;
    }

    public String[] getValidationOutputValues() {
        return validationOutputValues;
    }

    public String[][] getTestingInputValues() {
        return testingInputValues;
    }

    public String[] getTestingOutputValues() {
        return testingOutputValues;
    }
}
